"""
@Project    : cosmo-face
@Module     : config.py
@Author     : HuangJiWen[huangjiwen@haier.com]
@Created    : 2020/9/14 14:57
@Desc       : 训练配置文件
"""

from pathlib import Path

import torch
from easydict import EasyDict
from torch.nn import CrossEntropyLoss
from torchvision import transforms as trans


def get_config(training=True):
    conf = EasyDict()
    conf.data_path = Path('data')
    conf.work_path = Path('work_space/')
    # conf.model_path = conf.work_path/'irse50_pretrain_models' # pretrain
    conf.model_path = conf.work_path / 'irse50_models'
    conf.log_path = conf.work_path / 'log'
    conf.save_path = conf.work_path / 'save'
    conf.pretrain_path = conf.work_path / 'model'

    conf.input_size = [112, 112]
    conf.embedding_size = 512
    conf.use_pretrain = True
    conf.use_mobilfacenet = False
    conf.net_depth = 50
    conf.drop_ratio = 0.6
    conf.net_mode = 'ir_se'  # or 'ir'
    conf.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    conf.test_transform = trans.Compose([
        trans.ToTensor(),
        trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

    conf.data_mode = 'emore'
    conf.vgg_folder = conf.data_path / 'vgg_face_dataset/align'
    conf.ms1m_folder = conf.data_path / 'faces_ms1m_112x112'
    conf.emore_folder = conf.data_path / 'faces_emore'
    conf.batch_size = 100  # irse net depth 50
    # conf.batch_size = 200  # mobile facenet

    conf.pin_memory = True
    # conf.num_workers = 4  # when batch size is 200
    conf.num_workers = 3

    # --------------------Training Config ------------------------
    if training:
        conf.log_path = conf.work_path / 'log'
        conf.save_path = conf.work_path / 'save'
        # conf.weight_decay = 5e-4
        conf.lr = 1e-3
        conf.milestones = [12, 18, 21]
        # conf.milestones = [8, 12, 16]  # pretrain
        conf.momentum = 0.9
        conf.pin_memory = True
        # conf.num_workers = 4  # when batch size is 200
        conf.num_workers = 3
        conf.ce_loss = CrossEntropyLoss()
        # --------------------Inference Config ------------------------
    else:
        conf.face_database_path = conf.data_path / 'faces_database/align_data'
        conf.threshold = 1.5
        conf.face_limit = 10
        # when inference, at maximum detect 10 faces in one image, my laptop is slow
    return conf
